#!/usr/bin/env python
# -*- coding:utf-8 -*-
import sys

sys.path.append('..')
from common_utils import CBBCommonUtils


def get_plugin_info():
    plugin_info = {
        "name": "centos_intrusion_guard_005 Disable_ICMP_redirection_and_ICMP_source_routing",
        "plugin_id": "centos_intrusion_guard_005",
        "plugin_type": "Intrusion Guard",
        "info": "Disable ICMP redirection and source routing to prevent attackers from maliciously changing the system routing table with forged ICMP redirection messages.",
        "level": "B",
        "module": "Safety reinforcement",
        "author": "congshz",
        "keyword": "Safety reinforcement",
        "configable": "false"
    }
    return plugin_info


logger = None
cur_user = None
cur_module = None
cur_task = None


def set_plugin_logger(setter, user, module, task, *args, **kwargs):
    global logger, cur_user, cur_module, cur_task
    logger = setter
    cur_user = user
    cur_module = module
    cur_task = task


# 扫描函数
def scan(ip, sys_user, sys_pwd, flag=0):
    des_list = []
    error_count = 0
    # 获取该参数值
    cmd1 = "cat /proc/sys/net/ipv4/conf/all/accept_redirects"
    cmd2 = "cat /proc/sys/net/ipv4/conf/all/accept_source_route"
    result1, output1 = CBBCommonUtils.cbb_run_cmd(ip, cmd1, username=sys_user, passwd=sys_pwd)
    result2, output2 = CBBCommonUtils.cbb_run_cmd(ip, cmd2, username=sys_user, passwd=sys_pwd)
    # 检错，每次执行命令都会进行检错，并决定返回退出/继续执行
    if not result1 and not result2 and len(output1) != 0 and len(output2) != 0:
        # 结果检测及警告
        output1 = output1[0].strip()
        output2 = output2[0].strip()
        if output1 != "0" or output2 != "0":
            out_des = "DANGEROUS. ICMP_Accept_Redirects is {0}, ICMP_Source_route is {1}.".format(output1, output2)
            error_count += 1
            logger.debug_warning(cur_user, cur_module, cur_task+'_Scan', '', out_des)
            des_list.append(out_des)
        else:
            out_des = "SAFE. ICMP_Accept_Redirects and ICMP_Source_route are closed."
            des_list.append(out_des)
            logger.debug_info(cur_user, cur_module, cur_task+'_Scan', '', out_des)
    else:
        error_des = "CMD process ERROR. Error info: {0} {1}" .format(result1[0].replace("\n", ""), result2[0].replace("\n", ""))
        des_list.append(error_des)
        logger.debug_error(cur_user, cur_module, cur_task+'_Scan', '', error_des)
        error_count += 1
        return des_list, error_count
    logger.debug_info(cur_user, cur_module, cur_task+'_Scan', '', 'Scan Complete.')
    if error_count != 0:
        # 是否加固
        if flag == 1:
            reinforce_des, reinforce_err = reinforce(ip, sys_user, sys_pwd)
            des_list.extend(reinforce_des)
            if reinforce_err == 0:
                error_count = 0
    return des_list, error_count


# 加固函数
def reinforce(ip, sys_user, sys_pwd):
    des_list = []
    error_count = 0
    # 备份文件
    cmd1 = "cat /proc/sys/net/ipv4/conf/all/accept_redirects"
    cmd2 = "cat /proc/sys/net/ipv4/conf/all/accept_source_route"
    result1, output1 = CBBCommonUtils.cbb_run_cmd(ip, cmd1, username=sys_user, passwd=sys_pwd)
    result2, output2 = CBBCommonUtils.cbb_run_cmd(ip, cmd2, username=sys_user, passwd=sys_pwd)
    if not result1 and not result2 and len(output1) != 0 and len(output2) != 0:
        output1 = output1[0].strip()
        output2 = output2[0].strip()
        create_cmd1 = "echo " + output1[0].replace("\n", "") + "　> /etc/accept_redirects.`date +%Y-%m-%d`"
        create_cmd2 = "echo " + output2[0].replace("\n", "") + "　> /etc/accept_source_route.`date +%Y-%m-%d`"
        result1, output1 = CBBCommonUtils.cbb_run_cmd(ip, create_cmd1, username=sys_user, passwd=sys_pwd)
        result2, output2 = CBBCommonUtils.cbb_run_cmd(ip, create_cmd2, username=sys_user, passwd=sys_pwd)
        if not result1 and not result2:
            # 检查是否备份成功
            check_cmd1 = "ls /etc | grep accept_redirects.`date +%Y-%m-%d`"
            check_cmd2 = "ls /etc | grep accept_source_route.`date +%Y-%m-%d`"
            check_result1, check_output1 = CBBCommonUtils.cbb_run_cmd(ip, check_cmd1, username=sys_user, passwd=sys_pwd)
            check_result2, check_output2 = CBBCommonUtils.cbb_run_cmd(ip, check_cmd2, username=sys_user, passwd=sys_pwd)
            if not check_result1 and not check_result2 and len(check_output1) != 0 and len(check_output2) != 0:
                backup_check_des = "Backup SUCCEED."
                logger.debug_info(cur_user, cur_module, cur_task+'_Reinforce', '', backup_check_des)
            else:
                error_count += 1
                backup_check_des = "Backup check FAILED."
                logger.debug_error(cur_user, cur_module, cur_task+'_Reinforce', '', backup_check_des)
                des_list.append(backup_check_des)
                return des_list, error_count
        else:
            error_count += 1
            backup_des = "Backup FAILED."
            logger.debug_error(cur_user, cur_module, cur_task+'_Reinforce', '', backup_des)
            des_list.append(backup_des)
            return des_list, error_count
    else:
        error_count += 1
        re_err = "CMD process ERROR. Error info: {0} {1}".format(result1[0].replace("\n", ""), result2[0].replace("\n", ""))
        logger.debug_error(cur_user, cur_module, cur_task+'_Reinforce', '', re_err)
        des_list.append(re_err)
        return des_list, error_count
    safe_cmd1 = "sysctl -w net.ipv4.conf.all.accept_redirects=0"
    safe_cmd2 = "sysctl -w net.ipv4.conf.default.accept_redirects=0"
    # 修改该值
    result1, output1 = CBBCommonUtils.cbb_run_cmd(ip, safe_cmd1, username=sys_user, passwd=sys_pwd)
    result2, output2 = CBBCommonUtils.cbb_run_cmd(ip, safe_cmd2, username=sys_user, passwd=sys_pwd)
    if not result1 and not result2:
        # 检查是否修改成功
        check_re_cmd1 = "cat /proc/sys/net/ipv4/conf/all/accept_redirects"
        check_re_cmd2 = "cat /proc/sys/net/ipv4/conf/all/accept_source_route"
        result1, output1 = CBBCommonUtils.cbb_run_cmd(ip, check_re_cmd1,  username=sys_user, passwd=sys_pwd)
        result2, output2 = CBBCommonUtils.cbb_run_cmd(ip, check_re_cmd2,  username=sys_user, passwd=sys_pwd)
        if not result1 and not result2 and len(output1) != 0 and len(output2) != 0:
            output1 = output1[0].strip()
            output2 = output2[0].strip()
            if output1 != "0" or output2 != "0":
                out_des = "DANGEROUS. ICMP_Accept_Redirects is {0}, ICMP_Source_route is {1}.".format(output1, output2)
                error_count += 1
                logger.debug_warning(cur_user, cur_module, cur_task+'_Reinforce', '', out_des)
                des_list.append(out_des)
                des = "Reinforce FAILED."
                logger.debug_info(cur_user, cur_module, cur_task+'_Reinforce', '', des)
                des_list.append(des)
                # 加固失败，回滚
                rollback_des, rollback_err = rollback(ip, sys_user, sys_pwd)
                des_list.extend(rollback_des)
                return des_list, error_count
            else:
                out_des = "SAFE. ICMP_Accept_Redirects and ICMP_Source_route are closed."
                logger.debug_info(cur_user, cur_module, cur_task+'_Reinforce', '', out_des)
                des_list.append(out_des)
        else:
            error_count += 1
            re_err = "CMD process ERROR. Error info: {0} {1}" .format(result1[0].replace("\n", ""), result2[0].replace("\n", ""))
            logger.debug_error(cur_user, cur_module, cur_task+'_Reinforce', '', re_err)
            des_list.append(re_err)
            return des_list, error_count
    else:
        error_count += 1
        des = "Reinforce FAILED."
        logger.debug_error(cur_user, cur_module, cur_task+'_Reinforce', '', des)
        des_list.append(des)
        # 加固失败，回滚
        rollback_des, rollback_err = rollback(ip, sys_user, sys_pwd)
        des_list.extend(rollback_des)
        return des_list, error_count
    des_list.append("Reinforce Complete.")
    logger.debug_info(cur_user, cur_module, cur_task+'_Reinforce', '', "Reinforce Complete.")
    return des_list, error_count


# 回滚函数
def rollback(ip, sys_user, sys_pwd):
    des_list = []
    error_count = 0
    # 查找备份文件，此处的查找支持多备份文件的查找（结果会按时间顺序排序），并默认回滚最新一次备份的内容，可以连续多次回滚
    grep_cmd1 = "ls /etc/ | grep accept_redirects.*"
    grep_cmd2 = "ls /etc/ | grep accept_source_route.*"
    result1, output1 = CBBCommonUtils.cbb_run_cmd(ip, grep_cmd1, username=sys_user, passwd=sys_pwd)
    result2, output2 = CBBCommonUtils.cbb_run_cmd(ip, grep_cmd2, username=sys_user, passwd=sys_pwd)
    if not result1 and not result2 and len(output1) != 0 and len(output2) != 0:
        logger.debug_info(cur_user, cur_module, cur_task+'_Rollback', '', "Backup file FOUND.")
        target1 = output1[len(output1) - 1].replace("\n", "")
        target2 = output2[len(output2) - 1].replace("\n", "")
        cmd1 = "cat /etc/" + target1
        cmd2 = "cat /etc/" + target2
        result1, output1 = CBBCommonUtils.cbb_run_cmd(ip, cmd1, username=sys_user, passwd=sys_pwd)
        result2, output2 = CBBCommonUtils.cbb_run_cmd(ip, cmd2, username=sys_user, passwd=sys_pwd)
        if not result1 and not result2 and len(output1) != 0 and len(output2) != 0:
            # 回滚文件
            re_cmd1 = "echo " + str(output1[0])[0] + " > /proc/sys/net/ipv4/conf/all/accept_redirects"
            re_cmd2 = "echo " + str(output2[0])[0] + " > /proc/sys/net/ipv4/conf/all/accept_source_route"
            result1, output1 = CBBCommonUtils.cbb_run_cmd(ip, re_cmd1, username=sys_user, passwd=sys_pwd)
            result2, output2 = CBBCommonUtils.cbb_run_cmd(ip, re_cmd2, username=sys_user, passwd=sys_pwd)
            if not result1 and not result2:
                logger.debug_info(cur_user, cur_module, cur_task+'_Rollback', '', "Rollback SUCCEED.")
                del_cmd1 = "rm -rf /etc/" + target1
                del_cmd2 = "rm -rf /etc/" + target2
                # 删除对应的备份文件
                result1, output1 = CBBCommonUtils.cbb_run_cmd(ip, del_cmd1, username=sys_user, passwd=sys_pwd)
                result2, output2 = CBBCommonUtils.cbb_run_cmd(ip, del_cmd2, username=sys_user, passwd=sys_pwd)
                if not result1 and not result2:
                    del_des1 = "Backup file deleted. Target is {}.".format(target1)
                    logger.debug_info(cur_user, cur_module, cur_task+'_Rollback', '', del_des1)
                    del_des2 = "Backup file deleted. Target is {}.".format(target2)
                    logger.debug_info(cur_user, cur_module, cur_task+'_Rollback', '', del_des2)
                else:
                    error_count += 1
                    del_des = "Backup file delete FAILED."
                    des_list.append(del_des)
                    logger.debug_error(cur_user, cur_module, cur_task+'_Rollback', '', del_des)
                    return des_list, error_count
            else:
                error_count += 1
                des = "Rollback FAILED, please retry."
                logger.debug_error(cur_user, cur_module, cur_task+'_Rollback', '', des)
                des_list.append(des)
                return des_list, error_count
        else:
            error_count += 1
            des = "READ Rollback File FAILED, please retry."
            logger.debug_error(cur_user, cur_module, cur_task+'_Rollback', '', des)
            des_list.append(des)
            return des_list, error_count
    else:
        error_count = -1
        des_list.append("Backup file NOT FOUND.")
        logger.debug_error(cur_user, cur_module, cur_task+'_Rollback', '', "Backup file not found.")
        return des_list, error_count
    des_list.append("Rollback Complete.")
    logger.debug_info(cur_user, cur_module, cur_task+'_Rollback', '', "Rollback Complete.")
    return des_list, error_count


def check(ip, *args, **kwargs):
    sys_user = kwargs.get("system_user")
    sys_pwd = kwargs.get("system_pwd")
    comm = kwargs.get("command")
    try:
        des_list = []
        des_start = "centos_intrusion_guard_005 Start."
        logger.debug_info(cur_user, cur_module, cur_task+'_Check', '', des_start)
        if comm == 1:
            des_scan, error_scan = scan(ip, sys_user, sys_pwd, flag=0)
            des_list.extend(des_scan)
            step_error = int(error_scan)
        elif comm == 2:
            des_reinforce, error_reinforce = scan(ip, sys_user, sys_pwd, flag=1)
            des_list.extend(des_reinforce)
            step_error = int(error_reinforce)
        elif comm == 3:
            des_rollback, error_rollback = rollback(ip, sys_user, sys_pwd)
            des_list.extend(des_rollback)
            step_error = int(error_rollback)
        else:
            return {"code": 3, "count": 0, "des": ['command must be 1/2/3']}
        des_end = "centos_intrusion_guard_005 Complete."
        logger.debug_info(cur_user, cur_module, cur_task+'_Check', '', des_end)
        if step_error == 0:
            code = 0
        elif step_error <= -1:
            code = 2
        else:
            code = 1
        return {"code": code, "count": step_error, "des": des_list}
    except Exception as er:
        code = 1
        des = ["ERROR:", str(er)]
        logger.debug_error(cur_user, cur_module, cur_task+'_Check', '', des)
        return {"code": code, "count": 0, "des": des}


# check(ip="127.0.0.1", system_user="root", system_pwd="1qaz!QAZ", command=0, flag=0)
